#include <cstdio>
#include <stdint.h>

#include <iostream>
#include <set>
#include <string>
#include <sstream>
#include <queue>

#include <pthread.h>
#include <unistd.h>

#include "SNMPSession.h"
#include "Network.h"
#include "port.h"
#include "AgentProto.h"
#include "Resolver.h"

#include "keep_alive/ClientAlive.h"
#include "IListener.h"
#include "Utils.h"

// ================================================================================

static std::string agentId;
static CachingResolver nameResolver;

// --------------------------------------------------------------------------------

static IPAddress fromOid(const oid* s, int pos) {
  return s[pos] | s[pos + 1] << 8 | s[pos + 2] << 16 | s[pos + 3] << 24;
}

// --------------------------------------------------------------------------------

class Worker {
public:

  /**
   *  Initializes the agent
   *
   *  @param root  the IP address of the router to start network discovery from
   */

  Worker(char* root) {
    in_addr addr;

    inet_aton(root, &addr);
    _root = addr.s_addr;

    pthread_spin_init(&_rlock, 0);

    //pthread_create(&_worker, NULL, workerThread, this);
    pthread_create(&_resolver, NULL, resolverThread, this);
    _network.id = agentId;
  }

  ~Worker() {
    pthread_spin_destroy(&_rlock);
  }

  std::string serialize() {
    std::stringstream stream;
    text_oarchive ar(stream);

    lock();
    ar << _network;
    unlock();

    return stream.str();
  }


  /**
   *  The main loop wrapper for a non-forking worker
   */

  void mainLoop() {
    workerThread(this);
  }

private:
  void lock() {
    pthread_spin_lock(&_rlock);
  }

  void unlock() {
    pthread_spin_unlock(&_rlock);
  }

  bool visited(IPAddress ipAddr) {
    return _visited.find(ipAddr) != _visited.end();
  }


  /**
   *  Prepares SNMP information fetching
   */

  void initFetch() {
    _visited.clear();
    _q.push(_root);

    lock();
    _network.clear();
    unlock();
  }


  /**
   *  Starts SNMP information fetching
   *
   *  @param node  the IP address of the node to fetch information from
   */

  void fetchAll(IPAddress node) {
    SNMPSession ssn(toString(node));
    _network.nodes.push_back(Network::Node());

    _curNode = &_network.nodes.back();

    // Extract system name

    ssn.start(".1.3.6.1.2.1.1.1");
    netsnmp_variable_list* var = ssn.next();
    if (var) {
      _curNode->descr = std::string((const char*)var->val.string);
    }    

    std::cout << "Fetching the interface table..." << std::endl;
    fetchIfTable(&ssn);

    if (_curNode->ifs.empty()) {
      std::cout << "No interfaces found, giving up the node" << std::endl;

      _network.nodes.pop_back();
      return;
    }

    std::cout << "Fetching the client table..." << std::endl;
    fetchATTable(&ssn);

    std::cout << "Fetching the load table..." << std::endl;
    fetchLoad(&ssn);
  }


  /**
   *  Fetches a client table
   */

  void fetchATTable(SNMPSession* s) {
    /*
     *  This SNMP table maps interface indexes into IP addresses 
     *  this interface knows about
     */

    s->start(".1.3.6.1.2.1.3.1.1.3"); // iso(1) org(3) dod(6) internet(1) mgmt(2) mib-2(1) at(3) atTable(1) atEntry(1) atNetAddress(3)

    for (;;) {
      netsnmp_variable_list* var = s->next();
      if (!var) {
        break;
      }

      lock();

      IPAddress ipAddr = (IPAddress)var->val.objid[0];
      
      if (!visited(ipAddr)) {
        _visited.insert(ipAddr);
        _q.push(ipAddr);
        
        std::cout << "Found node " << toString(ipAddr) << std::endl;

        _network.addInterface(ipAddr);
      }

      unlock();
    }
  }


  /**
   *  Fetches an interface table
   */

  void fetchIfTable(SNMPSession* s) {
    s->start(".1.3.6.1.2.1.4.20.1.2"); // iso(1) org(3) dod(6) internet(1) mgmt(2) mib-2(1) ip(4) ipAddrTable(1) ipAddrEntry(1) ipAdEntIfIndex(2)

    //  Collect interface indexes 

    for (;;) {
      netsnmp_variable_list* var = s->next();
      if (!var) {
        break;
      }

      IPAddress ipAddr = fromOid(var->name, 10);
      
      if (ipAddr != INADDR_LOOPBACK) { // Ignore a loopback interface
        unsigned int ifIdx = *var->val.integer;

        _network.interfaces.insert(std::make_pair(ipAddr, Network::Interface(ipAddr)));
        Network::Interface* iface = _network.interface(ipAddr);

        _curNode->addInterface(ifIdx, ipAddr);

        _visited.insert(ipAddr);       // Prevent the node from being visited once again
      }
    }

    
    // Collect iterface subnet masks

    s->start(".1.3.6.1.2.1.4.20.1.3"); // iso(1) org(3) dod(6) internet(1) mgmt(2) mib-2(1) ip(4) ipAddrTable(1) ipAddrEntry(1) ipEntAdNetMask(3)

    for (;;) {
      netsnmp_variable_list* var = s->next();
      if (!var) {
        break;
      }

      IPAddress ipAddr = fromOid(var->name, 10);
      if (ipAddr != INADDR_LOOPBACK) {
        IPAddress netmask = (IPAddress)var->val.objid[0];
        Network::Interface* iface = _network.interface(ipAddr);

        iface->netmask = netmask;
        _network.addSubnet(iface);
      }
    }
  }

  /**
   *  Fetches load information
   */

  void fetchLoad(SNMPSession* s) {
    s->start(".1.3.6.1.2.1.2.2.1.10");  // iso(1) org(3) dod(6) internet(1) mgmt(2) mib-2(1) interfaces(2) ifTable(2) ifEntry(1) ifInOctets

    for (;;) {
      netsnmp_variable_list* var = s->next();
      
      if (!var) {
        break;
      }

      size_t ifIdx = var->name[10];
      Network::Node::IfIndexes::iterator it = _curNode->ifIndexes.find(ifIdx);
      if (it != _curNode->ifIndexes.end()) {
        IPAddress addr = it->second;
        _network.interface(addr)->outBytes = *var->val.integer;
      }
    }
    

    s->start(".1.3.6.1.2.1.2.2.1.16");  // iso(1) org(3) dod(6) internet(1) mgmt(2) mib-2(1) interfaces(2) ifTable(2) ifEntry(1) ifInOctets
    for (;;) {
      netsnmp_variable_list* var = s->next();
      
      if (!var) {
        break;
      }

      size_t ifIdx = var->name[10];
      Network::Node::IfIndexes::iterator it = _curNode->ifIndexes.find(ifIdx);
      if (it != _curNode->ifIndexes.end()) {
        IPAddress addr = it->second;
        _network.interface(addr)->inBytes = *var->val.integer;
      }
    }
  }


  /**
   *  The main agent procedure
   */
  
  static void* workerThread(void* argv) {
    Worker* self = (Worker*)argv;

    self->initFetch();

    for (;;) {
      std::cout << "Restarting router walk..." << std::endl;

      while (!self->_q.empty()) {
        IPAddress node = self->_q.front(); self->_q.pop();
        std::cout << "Examining " << toString(node) << std::endl;

        self->fetchAll(node);

        sleep(1);
      }

      self->initFetch();
    }

    return NULL;
  }

  
  /**
   *  The reverse resolver thread
   */

  typedef std::deque<IPAddress> IfAddresses;
    
  void resolver() {
    Worker* self = this;

    for (;;) {
      IfAddresses addrs = self->getInterfaceAddresses();

      for (IfAddresses::const_iterator it = addrs.begin();
           it != addrs.end();
           ++it) {
        Network::Interface* iface = self->_network.interface(*it);

        if (!iface) {
          std::cout << "Resolving " << toString(*it) << std::endl;
          std::string name = nameResolver.resolve(*it);
        } else if (iface->name.empty()) {
          std::cout << "Resolving " << toString(*it) << std::endl;
          std::string name = nameResolver.resolve(*it);

          self->lock();
          
          if (iface) {
            iface->name = name;
          }
          
          self->unlock();
        }

        sleep(2);
      }
    }
  }

  static void* resolverThread(void* arg) {
    Worker* self = (Worker*)arg;
    self->resolver();
    return NULL;
  }


  /**
   *  Returns a list of currently discovered interfaces
   */

  IfAddresses getInterfaceAddresses() {
    IfAddresses res;

    lock();

    for (Network::Interfaces::const_iterator it = _network.interfaces.begin();
         it != _network.interfaces.end();
         ++it) {
      res.push_back(it->second.address);
    }
    
    unlock();

    return res;
  }

private:
  /* Signle initialized state */

  IPAddress _root;
  pthread_t _worker, _resolver;
  pthread_spinlock_t _rlock;  // Synchronizes network data

  /* Discovery state */

  std::queue<IPAddress> _q;
  std::set<IPAddress> _visited;
  Network::Node* _curNode;

  Network _network;
};

// --------------------------------------------------------------------------------

class AgentReplier : public net::IListener {
public:
  AgentReplier(int port, Worker* worker, const char* collectorHost) { 
    _client.subscribe(this);

    for(;;) {
      try {
        std::cout << "Connecting to the collector..." << std::endl;
        _client.connectServer(collectorHost, port);
      } catch (net::NetException& e) {
        std::cout << "A net exception was caught while connecting to the collector server: " << e.what() << std::endl;
        std::cout << "Retrying after 10 sec..." << std::endl;
        sleep(10);
        continue;
      }

      break;
    } 

    _worker = worker;
  }

  virtual ~AgentReplier() { }

  std::vector<char> onRawReceived(const std::vector<char>& rx) {
    AgentRequest* msg = (AgentRequest*)&rx[0];

    switch (msg->type) {
    case AgentRequest::AIM_GETGRAPH: {
      std::string s = _worker->serialize();
      
      std::cout << "Sending network graph..." << std::endl;
      return stringToVector(s);       
    }

    case AgentRequest::AIM_GETID:
      return stringToVector(agentId);
      break;

    default:
      std::cout << "Invalid inbound message of type " << msg->type << std::endl;
    }

    return std::vector<char>();
  }

private:
  net::ClientAlive _client;
  Worker* _worker;
};

// --------------------------------------------------------------------------------

int main(int argc, char* argv[]) {
  int port = NGU_AGENT_PORT;
  char* target;
  char* collector;

  if (argc < 4) {
    printf("Usage: agent <Network node IP address> <Network ID> <Collector host> [<Collector port>] \n");
    return 1;
  }

  target = argv[1];
  agentId = argv[2];
  collector = argv[3];

  if (argc > 4) {
    port = atoi(argv[4]);
  }

  Worker worker(target);
  AgentReplier replier(port, &worker, collector);
  
  worker.mainLoop();
               
  return 0;
}
  
